from rtl_token import *
from circuit import *
from exception import RTLParseException

class RTLParser():

    SYNTAX_CHECK = False
    
    def __init__(self, tokens : list[Token]) -> None:
        self.__tokens : list[Token] = tokens
        self.__current_index : int = 0

        self.__module_name : str = None

        self.__input_wires : dict[str, dict[int, Wire] | Wire] = {}
        self.__output_wires : dict[str, dict[int, Wire] | Wire] = {}
        self.__internal_wires : dict[str, dict[int, Wire] | Wire] = {}
    
        self.__wires : set[Wire] = set()
        self.__circuits : set[Circuit] = set()

    def parse(self, module_type):
        self.__parse_module_header()

        self.__parse_ports()

        self.__parse_wires()

        self.__parse_circuits()

        self.__ensure_tokens('A module must end with keyword \'endmodule\'', TokenType.ENDMODULE)

        return module_type(self.__module_name, self.__input_wires, self.__output_wires, self.__wires, self.__circuits)

    def __parse_module_header(self):
        """
            module module_name (port_name1, port_name2, ... port_namen);
        """
        # First token must be a 'module'
        self.__ensure_tokens('A module must begin with keyword \'module\'', TokenType.MODULE)

        # Second token is module name
        self.__module_name = self.__ensure_tokens('A module name should be a \'IDENTIFIER\'', TokenType.IDENTIFIER)

        # Ignore the parament list
        while self.__advance_token().token_type != TokenType.SEMICOLON:
            pass

    def __parse_ports(self):
        while self.__check_tokens(TokenType.INPUT, TokenType.OUTPUT):
            self.__parse_port()

    def __parse_port(self):
        """
            input/output port_name, port_name, ... port_name;
        """
        token = self.__advance_token()
        wires = self.__input_wires if token.token_type == TokenType.INPUT else self.__output_wires
    
        # At least one port_name
        self.__parse_signal_name(wires)
        # Loop for each prot_name
        while self.__check_tokens(TokenType.COMMA):
            self.__advance()
            self.__parse_signal_name(wires)
        
        self.__ensure_tokens('Except \';\'', TokenType.SEMICOLON)

    def __parse_wires(self):
        while self.__match_tokens(TokenType.WIRE):
            self.__parse_wire()

    def __parse_wire(self):
        """
            wire w1, w2, ... , wn;
        """
        self.__parse_signal_name(self.__internal_wires)
        while self.__match_tokens(TokenType.COMMA):
            self.__parse_signal_name(self.__internal_wires)

        self.__ensure_tokens('Except \';\' at the end od wire declaration', TokenType.SEMICOLON)

    def __parse_circuits(self):
        while self.__check_tokens(TokenType.IDENTIFIER):
            self.__parse_circuit()

    def __parse_circuit(self):
        """
            CIRCUIT_TYPE name (.XX(xx), .YY(yy), ... .ZZ(zz))
        """
        type = self.__advance_token().literal
        name = self.__advance_token().literal

        self.__ensure_tokens('Argument list should begin with \'(\'', TokenType.LEFT_PAREN)
        
        arguments : dict[str, Wire] = {}
        self.__parse_argument(arguments)
        while self.__match_tokens(TokenType.COMMA):
            self.__parse_argument(arguments)
        
        self.__ensure_tokens('Argument list should end with \')\'', TokenType.RIGHT_PAREN)
        self.__ensure_tokens( 'Except \';\' at the end od circuit init', TokenType.SEMICOLON)

        self.__circuits.add(CircuitFactory.create_circuit(type, name, arguments))

    def __parse_argument(self, arguments : dict[str, Wire]):
        """
            .XX(xx)
        """
        self.__ensure_tokens('Argument should start with \'.\'', TokenType.DOT)
        port_name : str = \
            self.__ensure_tokens('A prot\'s name should be \'IDENTIFIER\'', TokenType.IDENTIFIER)
        self.__ensure_tokens('Except \'(\'', TokenType.LEFT_PAREN)

        signal_name : str = \
            self.__ensure_tokens('A signal\'s name should be \'IDENTIFIER\'', TokenType.IDENTIFIER)
        signal = self.__get_signal_wire(signal_name)

        if self.__check_tokens(TokenType.LEFT_BRACKET):
            self.__advance()
            index = self.__ensure_tokens('A index should be \'INTEGER\'', TokenType.INTEGER)
            self.__ensure_tokens('Except \']\'', TokenType.RIGHT_BRACKET)
            wire = signal[index]
        else:
            wire = signal

        self.__ensure_tokens('Except \')\'', TokenType.RIGHT_PAREN)

        arguments[port_name] = wire

    def __parse_signal_name(self, wires : dict[str, dict[int, Wire] | Wire]):
        if self.__check_tokens(TokenType.IDENTIFIER):
            signal_name = self.__advance_token().literal
            w = Wire()
            self.__wires.add(w)
            wires[signal_name] = w
        else:
            # [xx:xx] name
            self.__ensure_tokens('except \'[\'', TokenType.LEFT_BRACKET) 
            left_index = self.__ensure_tokens('A index should be \'INTEGER\'', TokenType.INTEGER)
            self.__ensure_tokens('except \':\' between [xx:xx]', TokenType.COLON)
            right_index = self.__ensure_tokens('A index should be \'INTEGER\'', TokenType.INTEGER)
            self.__ensure_tokens('except \']\'', TokenType.RIGHT_BRACKET)

            interval = -1 if left_index > right_index else 1
            wires_dict = { i : Wire() for i in range(left_index, right_index - 1, interval)}
            for w in wires_dict.values():
                self.__wires.add(w)

            wires[self.__advance_token().literal] = wires_dict

    def __at_end(self) -> bool:
        return self.__tokens[self.__current_index].token_type == TokenType.EOF

    def __ensure_tokens(self, msg : str, *args) -> object:
        """
            Ensure weahter current token's type is one of args.
            If so, return current token's literal and advance, overwise rasise Exception.
        """
        if RTLParser.SYNTAX_CHECK == False:
            return self.__advance_token().literal
        else: 
            if self.__match_tokens(*args) == False:
                raise RTLParseException(msg)
            else:
                return self.__previous_token().literal

    def __check_tokens(self, *args) -> bool:
        """
            Check weahter current token's type is one of args.
            If so, return True, overwise return False.        
        """
        token_type = self.__tokens[self.__current_index].token_type
        return (token_type in args)

    def __match_tokens(self, *args) -> bool:
        """
            Check weahter current token's type is one of args.
            If so, return True and advance current token, overwise return False and do noting.
        """
        token_type = self.__tokens[self.__current_index].token_type
        if token_type in args:
            self.__current_index += 1
            return True
        return False
    
    def __advance(self) -> None:
        """
            Advance current token
        """
        self.__current_index += 1

    def __advance_token(self) -> Token:
        """
            Return and advance current token
        """
        token = self.__tokens[self.__current_index]
        self.__current_index += 1
        return token
    
    def __current_token(self) -> Token:
        """
            Retuen current token
        """
        return self.__tokens[self.__current_index]
    
    def __previous_token(self) -> Token:
        """
            Retuen previous token
        """
        return self.__tokens[self.__current_index]

    def __get_signal_wire(self, signal_name : str) -> Wire | dict[int, Wire]:
        # Find from internal wires
        wire = self.__internal_wires.get(signal_name)
        if wire != None:
            return wire
        
        # Find From input
        wire = self.__input_wires.get(signal_name)
        if wire != None:
            return wire
        
        # Find From output
        wire = self.__output_wires.get(signal_name)
        if wire != None:
            return wire
    